"""


作者：何逸凡
版本：1.0
"""

import math
import numpy as np
import zhplot
import matplotlib.pyplot as plt
from typing import List, Tuple
from sklearn.cluster import KMeans
import itertools
import time


# === 常量定义 ===
UAV_SPEED = 150  # m/s
UAV_DETECT_RANGE = 60000  # 60km
UAV_DETECT_ANGLE = 60  # 60度扇形
NUM_UAVS = 5  # 无人机数量

# 游戏区域关键点
A1 = np.array([0, 500000 + 100000 * math.sqrt(3)])
A2 = np.array([190000, 500000 + 100000 * math.sqrt(3)])
A3 = np.array([240000, 500000 + 50000 * math.sqrt(3)])
A4 = np.array([240000, 50000 * math.sqrt(3)])
A5 = np.array([190000, 0])
A6 = np.array([0, 0])
AOI = np.array([A2, A3, A4, A5])  # 目标区域

# 5个不同的起点（可根据实际需求调整）
w1 = np.array([0, 500000 + 100000 * math.sqrt(3)])
w2 = np.array([0, 400000 + 80000 * math.sqrt(3)])
w3 = np.array([0, 300000 + 60000 * math.sqrt(3)])
w4 = np.array([0, 200000 + 40000 * math.sqrt(3)])
w5 = np.array([0, 100000 + 20000 * math.sqrt(3)])
w_points = [w1, w2, w3, w4, w5]


def point_in_polygon(point, polygon):
    x, y = point
    n = len(polygon)
    inside = False
    p1x, p1y = polygon[0]
    for i in range(1, n + 1):
        p2x, p2y = polygon[i % n]
        if y > min(p1y, p2y):
            if y <= max(p1y, p2y):
                if x <= max(p1x, p2x):
                    if p1y != p2y:
                        xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
                    if p1x == p2x or x <= xinters:
                        inside = not inside
        p1x, p1y = p2x, p2y
    return inside

def generate_random_points_in_aoi(num_points=20):
    min_x = np.min(AOI[:, 0])
    max_x = np.max(AOI[:, 0])
    min_y = np.min(AOI[:, 1])
    max_y = np.max(AOI[:, 1])
    points = []
    while len(points) < num_points:
        x = np.random.uniform(min_x, max_x)
        y = np.random.uniform(min_y, max_y)
        point = np.array([x, y])
        if point_in_polygon(point, AOI):
            points.append(point)
    return np.array(points)

def cost_calulate(pre, current, next):
    """
    计算路径代价
    :param pre: 前一个点
    :param current: 当前点
    :param next: 下一个点
    :return: 代价值
    """
    dy = abs(next[1] - current[1])

    dx = next[0] - current[0]
    if dy <= 80000:
        return  -dx / 2
    if dx == 0:
        raise ValueError("dx为0，无法计算cost")
    if dx >= dy - 80000:
        cost = -dx / 2
    else:
        if (pre[1] - current[1]) * dy > 0:  # 同边
            cost = (dx ** 2 + (dy - 80000) ** 2) / (2 * dx) - dx
        else:
            cost = (dx ** 2 + dy ** 2) / (2 * dx) - dx
    return cost


def assign_points_to_trajectories(points: np.ndarray, w_points: List[np.ndarray], n_trajectories: int) -> Tuple[List[List[np.ndarray]], float]:
    # KMeans聚类
    kmeans = KMeans(n_clusters=n_trajectories, random_state=0).fit(points)
    labels = kmeans.labels_
    clusters = [[] for _ in range(n_trajectories)]
    for idx, label in enumerate(labels):
        clusters[label].append(points[idx])
    # 每组内贪心排序，起点为对应w点
    trajectories = []
    total_cost = 0
    for i in range(n_trajectories):
        cluster = clusters[i]
        start = w_points[i]
        unvisited = cluster.copy()
        traj = [start]
        curr = start
        pre = start
        cost = 0
        while unvisited:
            # 对每个候选点，计算cost_calulate(pre, curr, next)
            costs = [cost_calulate(pre, curr, pt) for pt in unvisited]
            min_idx = int(np.argmin(costs))
            next_pt = unvisited.pop(min_idx)
            cost += cost_calulate(pre, curr, next_pt)
            traj.append(next_pt)
            pre = curr
            curr = next_pt
        trajectories.append(traj)
        total_cost += cost
    return trajectories, total_cost

def brute_force_trajectory(points: list, w_start: np.ndarray) -> Tuple[list, float]:
    # points: List[np.ndarray]，w_start: np.ndarray
    # 只考虑x递增的排列
    if not points:
        return [w_start], 0
    # 先按x排序，保证递增
    points_sorted = sorted(points, key=lambda p: p[0])
    n = len(points_sorted)
    min_cost = float('inf')
    best_path = None
    for perm in itertools.permutations(points_sorted):
        valid = True
        for i in range(1, n):
            if perm[i][0] < perm[i-1][0]:
                valid = False
                break
        if not valid:
            continue
        traj = [w_start] + list(perm)
        cost = 0
        for i in range(2, len(traj)):
            # if i == 2:
            #     cost += cost_calulate(traj[i-2], traj[i-1], traj[i]) / 5  # W点到第一个点的cost除以5
            # else:
            cost += cost_calulate(traj[i-2], traj[i-1], traj[i])
        if cost < min_cost:
            min_cost = cost
            best_path = traj
    if best_path is None:
        best_path = [w_start] + points_sorted
        cost = 0
        for i in range(2, len(best_path)):
            if i == 2:
                cost += cost_calulate(best_path[i-2], best_path[i-1], best_path[i]) / 5
            else:
                cost += cost_calulate(best_path[i-2], best_path[i-1], best_path[i])
        min_cost = cost
    return best_path, min_cost

def total_cost(path, w_start):
    if not path:
        return 0
    cost = -np.linalg.norm(w_start - path[0])
    for i in range(1, len(path)):
        current = path[i-1]
        nextp = path[i]
        if i == 1:
            dx = nextp[0] - current[0]
            dy = abs(nextp[1] - current[1])
            cost += (dx ** 2 + (dy - 80000) ** 2) / (2 * dx)
        else:
            pre = path[i-2]
            cost += cost_calulate(pre, current, nextp)
    return cost

def visualize_trajectories(trajectories: List[List[np.ndarray]], w_points: List[np.ndarray]):
    plt.figure(figsize=(12, 10))
    # 绘制AOI区域
    aoi_closed = np.vstack([AOI, AOI[0]])
    plt.plot(aoi_closed[:, 0], aoi_closed[:, 1], 'b-', linewidth=2, label='AOI区域')
    plt.fill(aoi_closed[:, 0], aoi_closed[:, 1], alpha=0.3, color='lightblue')
    # 绘制起点
    for i, w in enumerate(w_points):
        plt.plot(w[0], w[1], 'go', markersize=10)
        plt.annotate(f'W{i+1}', (w[0], w[1]), xytext=(-15, 10), textcoords='offset points', fontsize=10, fontweight='bold')
    # 绘制轨迹
    colors = plt.cm.get_cmap('tab10', len(trajectories))
    for i, traj in enumerate(trajectories):
        traj_arr = np.array(traj)
        plt.plot(traj_arr[:, 0], traj_arr[:, 1], '-', color=colors(i), linewidth=2, label=f'轨迹{i+1}')
        plt.scatter(traj_arr[1:, 0], traj_arr[1:, 1], color=colors(i), s=40)
        for j, pt in enumerate(traj_arr[1:]):
            plt.annotate(f'B{j+1}', (pt[0], pt[1]), xytext=(5, 5), textcoords='offset points', fontsize=8)
    plt.xlabel('X坐标 (m)')
    plt.ylabel('Y坐标 (m)')
    plt.title('AOI区域内的轨迹分配')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()

def two_opt_path(path, w_start):
    """
    2-opt路径优化
    :param path: 当前路径
    :param w_start: 起点
    :return: 优化后的路径
    """
    n = len(path)
    if n < 2:
        return path

    best_path = path.copy()
    best_cost = total_cost(best_path, w_start)

    improved = True
    while improved:
        improved = False
        for i in range(1, n - 1):
            for j in range(i + 1, n):
                new_path = best_path[:i] + best_path[i:j+1][::-1] + best_path[j+1:]
                new_cost = total_cost(new_path, w_start)
                if new_cost < best_cost:
                    best_path = new_path
                    best_cost = new_cost
                    improved = True
    return best_path

def greedy_path(points: List[np.ndarray], w_start: np.ndarray) -> List[np.ndarray]:
    """
    贪心路径生成，保证x递增
    :param points: 候选点列表
    :param w_start: 起点
    :return: 贪心路径
    """
    if not points:
        return []
    # 先按x递增排序
    unvisited = sorted(points, key=lambda p: p[0])
    path = []
    curr = w_start
    while unvisited:
        # 只考虑x大于当前点的点
        candidates = [pt for pt in unvisited if pt[0] > curr[0]]
        if not candidates:
            break
        dists = [np.linalg.norm(curr - pt) for pt in candidates]
        min_idx = int(np.argmin(dists))
        next_point = candidates[min_idx]
        path.append(next_point)
        # 用索引删除
        for idx, pt in enumerate(unvisited):
            if np.allclose(pt, next_point):
                del unvisited[idx]
                break
        curr = next_point
    # 如果还有未访问点（即x未严格递增），丢弃
    return path

def cross_group_exchange(paths, w_points, max_iter=10):
    # paths: List[List[np.ndarray]]
    n = len(paths)
    for _ in range(max_iter):
        improved = False
        for i in range(n):
            for j in range(n):
                if i == j or not paths[i]:
                    continue
                for idx, pt in enumerate(paths[i]):
                    # 只允许插入到目标组x递增合法位置
                    target_group = paths[j]
                    insert_positions = [k for k in range(len(target_group)+1)
                        if (k==0 or pt[0]>target_group[k-1][0]) and (k==len(target_group) or pt[0]<target_group[k][0])]
                    for pos in insert_positions:
                        new_paths = [p.copy() for p in paths]
                        # 从i组删除
                        del new_paths[i][idx]
                        # 插入到j组合法位置
                        new_paths[j] = new_paths[j][:pos] + [pt] + new_paths[j][pos:]
                        # 组内2-opt（保证x递增）
                        for k in [i, j]:
                            if new_paths[k]:
                                new_paths[k] = sorted(new_paths[k], key=lambda p: p[0])
                        # 计算总代价
                        total = 0
                        for k in range(n):
                            if new_paths[k]:
                                total += total_cost(new_paths[k], w_points[k])
                        # 只要总代价下降就接受
                        old_total = 0
                        for k in range(n):
                            if paths[k]:
                                old_total += total_cost(paths[k], w_points[k])
                        if total < old_total:
                            paths = new_paths
                            improved = True
                            break
                    if improved:
                        break
                if improved:
                    break
            if improved:
                break
        if not improved:
            break
    return paths

def multi_kmeans_partition_with_w_assignment_and_exchange(points, w_points, n_clusters, n_init=10, exchange_iter=10):
    best_total_cost = float('inf')
    best_paths = None
    best_w_assign = None
    for seed in range(n_init):
        kmeans = KMeans(n_clusters=n_clusters, random_state=seed).fit(points)
        labels = kmeans.labels_
        clusters = [[] for _ in range(n_clusters)]
        for idx, label in enumerate(labels):
            clusters[label].append(points[idx])
        # 组内贪心+2-opt
        group_paths = []
        for i in range(n_clusters):
            if not clusters[i]:
                group_paths.append([])
                continue
            path = greedy_path(clusters[i], w_points[i])
            path = two_opt_path(path, w_points[i])
            group_paths.append(path)
        # 穷举W点与分组配对
        min_cost = float('inf')
        min_assign = None
        min_paths = None
        for perm in itertools.permutations(range(n_clusters)):
            total = 0
            for i, group_idx in enumerate(perm):
                if group_paths[group_idx]:
                    total += total_cost(group_paths[group_idx], w_points[i])
            if total < min_cost:
                min_cost = total
                min_assign = perm
                min_paths = [group_paths[group_idx] for group_idx in perm]
        # 跨组点交换+组内2-opt+重新配对
        min_paths = cross_group_exchange(min_paths, w_points, max_iter=exchange_iter)
        # 再次穷举配对
        min_cost2 = float('inf')
        min_assign2 = None
        min_paths2 = None
        for perm in itertools.permutations(range(n_clusters)):
            total = 0
            for i, group_idx in enumerate(perm):
                if min_paths[group_idx]:
                    total += total_cost(min_paths[group_idx], w_points[i])
            if total < min_cost2:
                min_cost2 = total
                min_assign2 = perm
                min_paths2 = [min_paths[group_idx] for group_idx in perm]
        if min_cost2 < best_total_cost:
            best_total_cost = min_cost2
            best_paths = min_paths2
            best_w_assign = min_assign2
    if best_paths is None:
        best_paths = [[] for _ in range(n_clusters)]
    return best_paths, best_total_cost, best_w_assign

def genetic_global_optimizer(points, w_points, num_uavs, time_limit=5.0, pop_size=50, elite_frac=0.2, mutation_rate=0.2):
    N = len(points)
    points = np.array(points)
    start_time = time.time()
    # 染色体编码：每个点分配到哪个无人机（0~num_uavs-1）
    def decode(ind):
        # ind: (N,) array, 每个元素为分组编号
        groups = [[] for _ in range(num_uavs)]
        for idx, g in enumerate(ind):
            groups[g].append(points[idx])
        # 组内顺序用贪心+2-opt
        paths = []
        for i in range(num_uavs):
            if groups[i]:
                path = greedy_path(groups[i], w_points[i])
                path = two_opt_path(path, w_points[i])
                paths.append(path)
            else:
                paths.append([])
        return paths
    def fitness(ind):
        paths = decode(ind)
        total = 0
        for i in range(num_uavs):
            if paths[i]:
                total += total_cost(paths[i], w_points[i])
        return total
    # 初始化种群
    pop = []
    for _ in range(pop_size):
        ind = np.random.randint(0, num_uavs, size=N)
        pop.append(ind)
    # 加入KMeans启发式个体
    kmeans = KMeans(n_clusters=num_uavs, random_state=0).fit(points)
    pop.append(kmeans.labels_)
    pop = pop[:pop_size]
    best_ind = None
    best_fit = float('inf')
    while time.time() - start_time < time_limit:
        fits = np.array([fitness(ind) for ind in pop])
        idx = np.argmin(fits)
        if fits[idx] < best_fit:
            best_fit = fits[idx]
            best_ind = pop[idx].copy()
        # 选择
        elite_num = max(1, int(elite_frac * pop_size))
        elite_idx = fits.argsort()[:elite_num]
        new_pop = [pop[i].copy() for i in elite_idx]
        # 交叉
        while len(new_pop) < pop_size:
            p1, p2 = np.random.choice(elite_idx, 2, replace=False)
            cut = np.random.randint(1, N)
            child = np.concatenate([pop[p1][:cut], pop[p2][cut:]])
            # 变异
            if np.random.rand() < mutation_rate:
                m_idx = np.random.randint(0, N)
                child[m_idx] = np.random.randint(0, num_uavs)
            new_pop.append(child)
        pop = new_pop
    # 最优解解码
    best_paths = decode(best_ind)
    return best_paths, best_fit

def group_points_by_y_with_cost(points, max_groups=5, cost_per_point=5000, max_cost=100000):
    """
    按y轴排序并分组，每组累计cost不超过max_cost，最多max_groups组。
    :param points: np.ndarray, shape=(N,2)
    :return: List[List[np.ndarray]]
    """
    points_sorted = sorted(points, key=lambda p: p[1])
    groups = []
    current_group = []
    current_cost = 0
    for pt in points_sorted:
        if len(groups) < max_groups - 1 and current_cost + cost_per_point > max_cost:
            groups.append(current_group)
            current_group = []
            current_cost = 0
        current_group.append(pt)
        current_cost += cost_per_point
    if current_group:
        groups.append(current_group)
    # 如果分组数不足max_groups，补空组
    while len(groups) < max_groups:
        groups.append([])
    return groups[:max_groups]

def group_points_by_y_kmeans(points, n_groups=5):
    """
    按y轴KMeans聚类分组。
    :param points: np.ndarray, shape=(N,2)
    :return: List[List[np.ndarray]]
    """
    from sklearn.cluster import KMeans
    y_values = np.array([p[1] for p in points]).reshape(-1, 1)
    kmeans = KMeans(n_clusters=n_groups, random_state=0).fit(y_values)
    labels = kmeans.labels_
    groups = [[] for _ in range(n_groups)]
    for idx, label in enumerate(labels):
        groups[label].append(points[idx])
    return groups

def optimal_group_w_assignment(groups, w_points):
    """
    对所有分组和w点做全排列配对，总代价为w点到组内第一个点的距离之和，返回最优配对和总代价。
    :param groups: List[List[np.ndarray]]
    :param w_points: List[np.ndarray]
    :return: (best_assignment, min_total_cost)
    """
    import itertools
    min_total_cost = float('inf')
    best_assignment = None
    n = len(groups)
    for perm in itertools.permutations(range(n)):
        total_cost = 0
        for i, group_idx in enumerate(perm):
            group = groups[group_idx]
            if group:
                # 组内x递增排序
                group_sorted = sorted(group, key=lambda p: p[0])
                dist = np.linalg.norm(w_points[i] - group_sorted[0])
                total_cost += dist
        if total_cost < min_total_cost:
            min_total_cost = total_cost
            best_assignment = perm
    return best_assignment, min_total_cost


def main():
    print("正在AOI区域内生成20个随机点...")
    random_points = generate_random_points_in_aoi(20)
    for idx, pt in enumerate(random_points):
        globals()[f"B{idx+1}"] = tuple(pt)
    print(f"成功生成 {len(random_points)} 个随机点:")
    for i, point in enumerate(random_points):
        print(f"B{i+1}: ({point[0]:.0f}, {point[1]:.0f})")

    # === 新方案：y轴KMeans聚类分组 + w点最优配对 ===
    groups = group_points_by_y_kmeans(random_points, n_groups=NUM_UAVS)
    # 先组内x递增排序
    groups_sorted = [sorted(group, key=lambda p: p[0]) for group in groups]
    best_assignment, min_total_cost = optimal_group_w_assignment(groups_sorted, w_points)
    if best_assignment is None:
        print("未找到有效的分组与w点配对方案！")
        return random_points
    print(f"\n[新方案] 最优配对: {best_assignment}, 总代价: {min_total_cost:.2f}")
    trajectories = []
    for i, group_idx in enumerate(best_assignment):
        group = groups_sorted[group_idx]
        if not group:
            continue
        traj = [w_points[i]] + group
        trajectories.append(traj)
        print(f"轨迹{i+1}（起点: {w_points[i]}，点数: {len(group)}）:")
        for j, pt in enumerate(group):
            print(f"  B{j+1}: ({pt[0]:.0f}, {pt[1]:.0f})")
        # === 新增：计算并打印每条轨迹的总cost（含每个b点+5000） ===
        cost = total_cost(group, w_points[i]) + 5000 * len(group)
        print(f"轨迹{i+1}的总cost: {cost:.2f}")
    # 可视化
    visualize_trajectories(trajectories, w_points)
    return random_points

if __name__ == "__main__":
    main()
